#!/bin/bash
# Author: <fanyuantao>
# Descripton: This script is used to set iptables rules for container communication.

function check_docker_is_exist() {
    if ! command -v docker >/dev/null 2>&1; then
        echo "docker is not installed, please install it first."
        exit 1
    fi

    if [ ! -d /etc/docker ]; then
        echo "docker config directory is not exist, create it."
        mkdir -p /etc/docker
    fi

    if [ -f /etc/docker/daemon.json ]; then
        cp /etc/docker/daemon.json /etc/docker/daemon.json.bak
        echo '{"log-driver": "json-file", "log-opts": {"max-size": "100m", "max-file": "3"}, "live-restore": true, "iptables": true}' > /etc/docker/daemon.json
    else
        echo '{"log-driver": "json-file", "log-opts": {"max-size": "100m", "max-file": "3"}, "live-restore": true, "iptables": true}' > /etc/docker/daemon.json
    fi
}

function open_iptables_for_docker() {
    if [ ! -z "$(grep -i '"iptables": false' /etc/docker/daemon.json)" ]; then
        sed -i "s/\"iptables\": false/\"iptables\": true/g" /etc/docker/daemon.json
        systemctl daemon-reload
        systemctl restart docker
    fi
    echo "success"
}

function close_iptables_for_docker() {
    if [ ! -z "$(grep -i '"iptables": true' /etc/docker/daemon.json)" ]; then
        sed -i "s/\"iptables\": true/\"iptables\": false/g" /etc/docker/daemon.json
        systemctl daemon-reload
        systemctl restart docker
    fi
    echo "success"
}

function list_iptables_for_docker() {
    echo "Current iptables rules :"
    iptables -nL
}

function look_config() {
    cat /etc/docker/daemon.json
}

function set_iptables_rules() {
    while true; do
        echo ""
        echo "1. Open iptables for docker"
        echo "2. Close iptables for docker"
        echo "3. List for docker iptables rules"
        echo "4. Look config"
        echo "0. Exit"
        echo ""
        read -p "Please select : " choice
        case $choice in
            1) open_iptables_for_docker;;
            2) close_iptables_for_docker;;
            3) list_iptables_for_docker;;
            4) look_config;;
            0) exit;;
            *) echo "Invalid choice, please try again.";;
        esac
    done
}

function main() {
    check_docker_is_exist
    set_iptables_rules
}   

main